#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2025/5/14 16:55
@Author  : HZP
@File    : 1.摘要缓冲混合记忆.py
"""
from typing import Any

from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()


# 1.max_tokens用于判断是否需要生成新的摘要
# 2.summary用于存储摘要的信息
# 3.chat_histories用于存储历史对话
# 4.get_num_tokens用于计算传入文本的token数
# 5.save_context用于存储新的交流对话
# 6.get_buffer_string用于将历史对话转换成字符串
# 7.load_memory_variables用于加载记忆变量信息
# 8.summary_text用于将旧的摘要和传入的对话生成新摘要
class ConversationSummaryBufferMemory:

    def __init__(self, max_tokens=30, summary=None, chat_histories=None):
        self.max_tokens = max_tokens
        self.summary = summary
        self.chat_histories = [] if chat_histories is None else chat_histories
        self.client = OpenAI(base_url="https://api-wc33k5m6j8f4ve1d.aistudio-app.com/v1")

    def summary_text(self, origin_summary: str, new_line: str) -> str:
        """生成新的摘要信息"""
        prompt = f"""你是一个强大的聊天机器人，请根据用户提供的谈话内容，总结摘要，并将其添加到先前提供的摘要中
        ，返回一个新的摘要，除了新摘要其他任何数据都不要生成，如果用户的对话信息里有一些关键的信息，比方说姓名、爱好、性别、重要事件等等，
        这些全部都要包括在生成的摘要中，摘要尽可能要还原用户的对话记录。
        请不要将<example>标签里的数据当成实际的数据，这里的数据只是一个示例数据，告诉你该如何生成新摘要。
        <example>
        当前摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量。
        新的对话：
        Human：为什么你认为人工智能是一股向善的力量？
        AI：因为人工智能会帮助人类充分发挥潜力。
        新摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量，因为它将帮助人类充分发挥潜力。
        </example>
        =====================以下的数据是实际需要处理的数据=====================
        当前摘要：{origin_summary}
        新的对话：
        {new_line}
        请帮用户将上面的信息生成新摘要。"""
        result_data = self.client.chat.completions.create(model="deepseek-r1:70b",
                                                          messages=[
                                                              {"role": "user", "content": prompt}],
                                                          stream=False)
        print(f"生成的摘要内容：{result_data.choices[0].message.content}")
        return result_data.choices[0].message.content

    @classmethod
    def get_num_tokens(cls, query: str) -> int:
        """计算传入文本的token数"""
        return len(query)

    def save_context(self, human_message: str, ai_message: str) -> None:
        """存储新的对话记录"""
        self.chat_histories.append({"human": human_message, "ai": ai_message})
        buffer = self.get_buffer_string()
        if self.get_num_tokens(buffer) > self.max_tokens:
            """进行摘要生成 并移除一个会话"""
            message = self.chat_histories[0]
            """生成摘要"""
            message_str = f"Human:{message.get('human')}\nAI:{message.get('ai')}"
            self.summary = self.summary_text(self.summary, message_str)
            del self.chat_histories[0]

    def get_buffer_string(self) -> str:
        """将历史对话转换成字符串"""
        buffer = ""
        for chat in self.chat_histories:
            buffer += f"Human:{chat['human']}\nAI:{chat['ai']}\n"
        return buffer.strip()

    def load_memory_variables(self) -> dict[str, Any]:
        """加载记忆变量信息"""
        history_message = self.get_buffer_string()
        return {"chat_history": f"摘要：{self.summary}\n\n历史信息:{history_message}"}


llm = OpenAI(base_url="https://api-wc33k5m6j8f4ve1d.aistudio-app.com/v1")

summary = ConversationSummaryBufferMemory()

while True:
    query = input("Human: ")
    # q 为退出会话
    if query == "q":
        break
    summary_message = summary.load_memory_variables()
    # 进行请求ai数据
    prompt = ("你是一个强大的聊天机器人，请根据对应的上下文和用户提问解决问题。\n\n"
              f"{summary_message.get('chat_history')}\n\n"
              f"用户的提问是:{query}")

    response = llm.chat.completions.create(
        model="deepseek-r1:70b",
        messages=[
            {"role": "user", "content": prompt}],
        stream=True
    )
    ai_content = ""
    print("AI: ", flush=True, end="")
    for content_item in response:
        content = content_item.choices[0].delta.content
        print(content, flush=True, end="")
        ai_content += content
    print("")
    summary.save_context(query, ai_content)
